import cv2
import numpy as np


def enhance_contrast(frame, method='clahe', clip_limit=2.0, grid_size=(8, 8)):
    """增强单帧图像的对比度"""
    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) if len(frame.shape) == 3 else frame

    if method == 'stretch':
        min_val = np.min(gray)
        max_val = np.max(gray)
        return np.uint8(255 * (gray - min_val) / (max_val - min_val))
    elif method == 'equalize':
        return cv2.equalizeHist(gray)
    elif method == 'clahe':
        clahe = cv2.createCLAHE(clipLimit=clip_limit, tileGridSize=grid_size)
        return clahe.apply(gray)
    return gray


def ncc_template_matching(src, template):
    """执行NCC模板匹配,用库函数输出结果当作ground truth"""
    res = cv2.matchTemplate(src, template, cv2.TM_CCOEFF_NORMED)
    _, max_val, _, max_loc = cv2.minMaxLoc(res)
    return max_val, max_loc


# 视频和模板路径
video_path = 'car.mp4'
template = cv2.imread('car_1.jpg', cv2.IMREAD_GRAYSCALE)

# 视频捕捉对象
cap = cv2.VideoCapture(video_path)

# 获取视频参数
frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = cap.get(cv2.CAP_PROP_FPS)

# 视频写入对象
fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # 编解码器
video_writer = cv2.VideoWriter('output.mp4', fourcc, fps, (frame_width, frame_height))

# 模板尺寸
w, h = template.shape[::-1]

while cap.isOpened():
    ret, frame = cap.read()
    if not ret:
        break

    # 对比度增强（返回灰度图）
    enhanced_gray = enhance_contrast(frame, method='clahe')

    # NCC模板匹配
    ncc_score, (x, y) = ncc_template_matching(enhanced_gray, template)

    # # 转换为三通道用于绘制彩色框
    result_frame = cv2.cvtColor(enhanced_gray, cv2.COLOR_GRAY2BGR)

    # 绘制匹配结果
    if ncc_score > 0.7:
        cv2.rectangle(result_frame, (x, y), (x + w, y + h), (0, 255, 0), 2)

    # 显示和保存
    cv2.imshow('Result', result_frame)
    video_writer.write(result_frame)

    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
video_writer.release()
cv2.destroyAllWindows()